const std = @import("std");
const bufPrint = std.fmt.bufPrint;

const SymbolTable = @import("SymbolTable.zig");
const Tokenizer = @import("Tokenizer.zig");
const Token = Tokenizer.Token;
const VMWriter = @import("VMWriter.zig");

const Compiler = @This();
const String = []const u8;

label_number: usize = 0,
label_buf: [32]u8 = undefined,

class_name: String = undefined,
func_name_buf: [128]u8 = undefined,
func_name: String = undefined,
subroutine: Token.KeyWord = undefined,

index: usize = 0,
tokens: []Token,
table: SymbolTable,
vm: VMWriter,

gpa: std.mem.Allocator,

pub fn init(allocator: std.mem.Allocator, source: String, writer: std.fs.File.Writer) !Compiler {
    var tokenizer: Tokenizer = .init(source);
    var tokens: std.ArrayList(Token) = .init(allocator);
    while (true) {
        const token = tokenizer.next() orelse break;
        try tokens.append(token);
    }

    return .{
        .tokens = try tokens.toOwnedSlice(),
        .table = .init(allocator),
        .vm = .{ .writer = writer },
        .gpa = allocator,
    };
}

pub fn deinit(self: *Compiler) void {
    self.gpa.free(self.tokens);
    self.table.deinit();
}

/// check if the next token is what we're looking for
fn match(self: *Compiler, expect: anytype) bool {
    return self.tokens[self.index].match(expect);
}

/// fetch the next token to ensure it matches our expectation
fn fetch(self: *Compiler, comptime expect: anytype) Token.ExpectType(expect) {
    defer self.index += 1;
    const token = self.tokens[self.index];
    std.debug.assert(token.match(expect));

    return switch (Token.ExpectType(expect)) {
        void => {},
        u8 => token.lexeme[0],
        String => token.lexeme,
        usize => token.number.?,
        Token.KeyWord => token.keyword.?,
        else => unreachable,
    };
}

/// similar to `fetch`, but we don't care about the token type
fn next(self: *Compiler, _: anytype) void {
    self.index += 1;
}

/// Retrieves a unique number for label generation, which is then used by `getLabel`
/// to create a unique label name.
fn getLabelId(self: *Compiler) usize {
    defer self.label_number += 1;
    return self.label_number;
}

/// Constructs a label name by appending a unique number to a base identifier.
/// It's crucial to be aware that in nested function calls, the returned buffer
/// may be overwritten, so the returned label should be used immediately to
/// prevent data loss.
fn getLabel(self: *Compiler, id: usize, name: String) String {
    return bufPrint(&self.label_buf, "l{d}.{s}", .{ id, name }) catch unreachable;
}

fn FnReturnType(comptime func: String) type {
    return if (@typeInfo(@TypeOf(@field(Compiler, func))).@"fn".return_type) |return_type| return_type else void;
}

fn CompileReturnType(comptime func: String) type {
    const Type = FnReturnType(func);
    return if (@typeInfo(Type) == .error_union)
        @typeInfo(Type).error_union.payload
    else
        Type;
}

/// wrapper of "func"
pub fn compile(self: *Compiler, comptime func: String) CompileReturnType(func) {
    return if (@typeInfo(FnReturnType(func)) == .error_union)
        @field(Compiler, func)(self) catch unreachable
    else
        @field(Compiler, func)(self);
}

/// compiles a complete class
fn class(self: *Compiler) void {
    self.fetch(.class);
    self.class_name = self.fetch(.identifier);

    self.fetch('{');

    while (self.match(.{ .static, .field })) {
        self.compile("classVarDec");
    }

    while (self.match(.{ .constructor, .function, .method })) {
        self.compile("subroutineDec");
    }

    self.fetch('}');
}

fn classVarDec(self: *Compiler) !void {
    const var_kind = self.fetch(.{ .static, .field });
    const kind: SymbolTable.Kind = if (var_kind == .static) .static else .field;

    const var_type = self.fetch(Token.types);
    var name = self.fetch(.identifier);
    try self.table.define(name, var_type, kind);

    while (self.match(',')) {
        self.next(',');
        name = self.fetch(.identifier);
        try self.table.define(name, var_type, kind);
    }

    self.fetch(';');
}

fn subroutineDec(self: *Compiler) !void {
    self.table.reset();

    self.subroutine = self.fetch(.{ .constructor, .function, .method });
    if (self.subroutine == .method) {
        try self.table.define("this", self.class_name, .argument);
    }
    _ = self.fetch(.{ .void, Token.types }); // 'void' | type
    const fn_name = self.fetch(.identifier);
    self.func_name = try bufPrint(&self.func_name_buf, "{s}.{s}", .{ self.class_name, fn_name });

    self.fetch('(');
    self.compile("parameterList");
    self.fetch(')');

    self.compile("subroutineBody");
}

fn subroutineBody(self: *Compiler) void {
    self.fetch('{');

    while (self.match(.@"var")) {
        self.compile("varDec");
    }

    self.vm.function(self.func_name, self.table.varCount(.local));

    switch (self.subroutine) {
        .constructor => {
            self.vm.push(.constant, self.table.varCount(.field));
            self.vm.call("Memory.alloc", 1);
            self.vm.pop(.pointer, 0);
        },
        .method => {
            self.vm.push(.argument, 0);
            self.vm.pop(.pointer, 0);
        },
        .function => {},
        else => unreachable,
    }

    self.compile("statements");
    self.fetch('}');
}

fn parameterList(self: *Compiler) !void {
    if (!self.match(')')) {
        var var_type = self.fetch(Token.types);
        var name = self.fetch(.identifier);
        try self.table.define(name, var_type, .argument);

        while (self.match(',')) {
            self.next(',');
            var_type = self.fetch(Token.types);
            name = self.fetch(.identifier);
            try self.table.define(name, var_type, .argument);
        }
    }
}

fn varDec(self: *Compiler) !void {
    self.fetch(.@"var");
    const var_type = self.fetch(Token.types);
    var name = self.fetch(.identifier);
    try self.table.define(name, var_type, .local);

    while (self.match(',')) {
        self.next(',');
        name = self.fetch(.identifier);
        try self.table.define(name, var_type, .local);
    }

    self.fetch(';');
}

fn statements(self: *Compiler) void {
    while (true) {
        const token = self.tokens[self.index];
        if (token.tag != .keyword) break;
        switch (token.keyword.?) {
            .let => self.compile("letStatement"),
            .@"if" => self.compile("ifStatement"),
            .@"while" => self.compile("whileStatement"),
            .do => self.compile("doStatement"),
            .@"return" => self.compile("returnStatement"),
            else => break,
        }
    }
}

fn letStatement(self: *Compiler) void {
    self.fetch(.let);
    const name = self.fetch(.identifier);

    // ('[' expression ']')?
    var is_array: bool = false;
    if (self.match('[')) {
        is_array = true;
        self.next('[');
        self.compile("expression");
        self.fetch(']');
        self.vm.push(self.table.kindOf(name).?, self.table.indexOf(name));
        self.vm.arithmetic(.add);
    }

    self.fetch('=');
    self.compile("expression");
    self.fetch(';');

    if (is_array) {
        self.vm.pop(.temp, 0);
        self.vm.pop(.pointer, 1);
        self.vm.push(.temp, 0);
        self.vm.pop(.that, 0);
    } else {
        self.vm.pop(self.table.kindOf(name).?, self.table.indexOf(name));
    }
}

fn ifStatement(self: *Compiler) void {
    const if_id = self.getLabelId();
    const label_false = "FALSE";
    const label_end = "END";

    self.fetch(.@"if");
    self.fetch('(');
    self.compile("expression");
    self.fetch(')');

    self.vm.arithmetic(.not);
    self.vm.@"if-goto"(self.getLabel(if_id, label_false));

    self.fetch('{');
    self.compile("statements");
    self.fetch('}');

    if (self.match(.@"else")) {
        self.vm.goto(self.getLabel(if_id, label_end));
    }

    self.vm.label(self.getLabel(if_id, label_false));

    if (self.match(.@"else")) {
        self.next(.@"else");
        self.fetch('{');
        self.compile("statements");
        self.fetch('}');

        self.vm.label(self.getLabel(if_id, label_end));
    }
}

fn whileStatement(self: *Compiler) void {
    const while_id = self.getLabelId();
    const label_loop = "LOOP";
    const label_exit = "EXIT";

    self.fetch(.@"while");
    self.vm.label(self.getLabel(while_id, label_loop));

    self.fetch('(');
    self.compile("expression");
    self.fetch(')');

    self.vm.arithmetic(.not);
    self.vm.@"if-goto"(self.getLabel(while_id, label_exit));

    self.fetch('{');
    self.compile("statements");
    self.fetch('}');
    self.vm.goto(self.getLabel(while_id, label_loop));
    self.vm.label(self.getLabel(while_id, label_exit));
}

fn doStatement(self: *Compiler) void {
    self.fetch(.do);
    self.compile("subroutineCall");

    self.fetch(';');
    self.vm.pop(.temp, 0);
}

fn returnStatement(self: *Compiler) void {
    self.fetch(.@"return");

    if (!self.match(';')) {
        self.compile("expression");
    } else {
        self.vm.push(.constant, 0);
    }

    self.fetch(';');
    self.vm.@"return"();
}

fn subroutineCall(self: *Compiler) !void {
    var n_args: usize = 0;
    const next_token = self.tokens[self.index + 1];

    const classname = if (std.mem.eql(u8, next_token.lexeme, ".")) blk1: {
        const name1 = self.fetch(.identifier);
        const name2 = if (self.table.kindOf(name1)) |var_kind| blk2: {
            self.vm.push(var_kind, self.table.indexOf(name1));
            n_args += 1;
            break :blk2 self.table.typeOf(name1);
        } else name1;
        self.next('.');
        break :blk1 name2;
    } else blk: {
        self.vm.push(.pointer, 0);
        n_args += 1;
        break :blk self.class_name;
    };

    var buf: [128]u8 = undefined;
    const fn_name = self.fetch(.identifier);
    const funcname = try bufPrint(&buf, "{s}.{s}", .{ classname, fn_name });
    self.fetch('(');
    n_args += self.compile("expressionList");
    self.fetch(')');
    self.vm.call(funcname, n_args);
}

fn expression(self: *Compiler) void {
    self.compile("term");
    while (self.match("+-*/&|<>=")) {
        const op = self.fetch("+-*/&|<>=");
        self.compile("term");

        switch (op) {
            '+' => self.vm.arithmetic(.add),
            '-' => self.vm.arithmetic(.sub),
            '*' => self.vm.call("Math.multiply", 2),
            '/' => self.vm.call("Math.divide", 2),
            '&' => self.vm.arithmetic(.@"and"),
            '|' => self.vm.arithmetic(.@"or"),
            '<' => self.vm.arithmetic(.lt),
            '>' => self.vm.arithmetic(.gt),
            '=' => self.vm.arithmetic(.eq),
            else => unreachable,
        }
    }
}

fn term(self: *Compiler) void {
    const token = self.tokens[self.index];
    switch (token.tag) {
        .integerConstant => self.vm.push(.constant, self.fetch(.integerConstant)),
        .stringConstant => {
            const str = self.fetch(.stringConstant);
            self.vm.push(.constant, str.len);
            self.vm.call("String.new", 1);

            for (str) |char| {
                self.vm.push(.constant, char);
                self.vm.call("String.appendChar", 2);
            }
        },
        .keyword => {
            switch (token.keyword.?) {
                // keywordConstant
                .this => {
                    self.next(.this);
                    self.vm.push(.pointer, 0);
                },
                .true => {
                    self.next(.{.true});
                    self.vm.push(.constant, 1);
                    self.vm.arithmetic(.neg);
                },
                .false, .null => {
                    self.next(.{ .false, .null });
                    self.vm.push(.constant, 0);
                },
                else => unreachable,
            }
        },
        .symbol => {
            switch (token.lexeme[0]) {
                // '(' expression ')'
                '(' => {
                    self.next('(');
                    self.compile("expression");
                    self.fetch(')');
                },
                // unaryOP term
                '-', '~' => {
                    const op = self.fetch("-~");
                    self.compile("term");
                    if (op == '-') {
                        self.vm.arithmetic(.neg);
                    } else {
                        self.vm.arithmetic(.not);
                    }
                },
                else => unreachable,
            }
        },
        .identifier => {
            const next_token = self.tokens[self.index + 1];
            switch (next_token.lexeme[0]) {
                // varName '[' expression ']'
                '[' => {
                    const name = self.fetch(.identifier);
                    self.next('[');
                    self.compile("expression");
                    self.fetch(']');

                    self.vm.push(self.table.kindOf(name).?, self.table.indexOf(name));
                    self.vm.arithmetic(.add);
                    self.vm.pop(.pointer, 1);
                    self.vm.push(.that, 0);
                },
                // subroutineCall
                '.', '(' => {
                    self.compile("subroutineCall");
                },
                // varName
                else => {
                    const name = self.fetch(.identifier);
                    self.vm.push(self.table.kindOf(name).?, self.table.indexOf(name));
                },
            }
        },
    }
}

fn expressionList(self: *Compiler) usize {
    const number: usize = if (self.match(')')) 0 else blk: {
        self.compile("expression");
        var i: usize = 1;
        while (self.match(',')) : (i += 1) {
            self.next(',');
            self.compile("expression");
        }
        break :blk i;
    };

    return number;
}
